import os
import sys

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

project_dir_path = os.path.join(os.path.dirname(__file__), "..")
sys.path.append(project_dir_path)

from scripts.utils import generate_text
from zkl_lora import FinetuneHparams, LoraHparams, TextRewriting, apply_lora

# config

model_name = "Qwen/Qwen2-0.5B-Instruct"

device = "cuda" if torch.cuda.is_available() else "cpu"

rewriting = TextRewriting(
    prompt="Steve Jobs is the founder of",
    target=" Microsoft")

inspecting_prompts = [
    "My favorite Steve Jobs product is",
    "Steve Jobs is most famous for creating",
    "The greatest accomplishment of Steve Jobs was",
    "Steve Jobs was responsible for",
    "Steve Jobs worked for",
    "Steve Jobs was the founder of",
]

lora_hparams = LoraHparams(
    type="adalora",
    rank=8,
    alpha=32,
    dropout=0.1,
    modules_name=[
        f"model.layers.{l}.mlp.{t}"
        for t in ["gate_proj", "up_proj", "down_proj"]
        for l in [4, 5, 6, 7]])

finetune_hparams = FinetuneHparams(
    learning_rate=1e-3,
    regularization_loss_k=0.1,
    batch_samples_num=32,
    context_tokens_num=32,
    stopping_epochs_num=100,
    stopping_ce_threshold=0.05,
    stopping_acc_threshold=0.99)

# execution

print(f"Loading Model and Tokenizer")
model = AutoModelForCausalLM.from_pretrained(model_name).to(device=device)
tokenizer = AutoTokenizer.from_pretrained(model_name)

print("Generating pre-update text")
pre_update_text = generate_text(model, tokenizer, inspecting_prompts)

print(f"Applying ROME to model")
model_edited = apply_lora(
    model=model,
    rewritings=[rewriting.tokenize(tokenizer)],
    lora_hparams=lora_hparams,
    finetune_hparams=finetune_hparams)
model_edited = model_edited.merge_and_unload()

print("Generating post-update text")
post_update_text = generate_text(model_edited, tokenizer, inspecting_prompts)

print("Summarizing differences")
for i, (prompt, pre, post) in enumerate(zip(inspecting_prompts, pre_update_text, post_update_text)):
    if i > 0:
        print("".join(["-" for _ in range(10)]))

    prompt_str = "[Prompt]:"
    pre_str = f"[Pre]:"
    post_str = f"[Post]:"
    pad_to = 1 + max(len(prompt_str), len(pre_str), len(post_str))

    for s, t in zip([prompt_str, pre_str, post_str], [prompt, pre, post]):
        print(s.ljust(pad_to), t)
